#include "pin.H"

VOID pinCheckpointerInstrumentTrace(TRACE trace, VOID *v)
{

  for( BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) ){

    for(INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) ){
      
      if( INS_IsStackWrite(ins) ){

        INS_InsertCall(ins, IPOINT_BEFORE,(AFUNPTR)UpdateStackSize, 
                       IARG_MEMORYWRITE_EA, IARG_MEMORYWRITE_SIZE, 
                       IARG_END);

      }

      if( INS_IsDirectCall(ins) && INS_IsProcedureCall(ins) ){

        ADDRINT add = INS_DirectBranchOrCallTargetAddress(ins);

        RTN rtn = RTN_FindByAddress(add);

        if( RTN_Valid(rtn) ){

          if(strstr(RTN_Name(rtn).c_str(),"LOAD_EXPLICIT_CHECKPOINT")){
              
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)RestoreCheckPoint, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }

          if(strstr(RTN_Name(rtn).c_str(),"STORE_EXPLICIT_CHECKPOINT")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TRY_BEGIN")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }
          if(strstr(RTN_Name(rtn).c_str(),"TRY_END")){
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)CaptureCheckPoint, 
                           IARG_CONTEXT, 
                           IARG_FUNCARG_CALLSITE_VALUE, 0, 
                           IARG_END);
          }

        }

      }

    }

  }

}

/*To be called in PIN main*/
void initCheckpointer()
{

  ckptr = new PinCheckpointer();
  TRACE_AddInstrumentFunction(pinCheckpointerInstrumentTrace, 0);

}
